package com.shujia.spark.core

import org.apache.spark.rdd.RDD
import org.apache.spark.{Partitioner, SparkConf, SparkContext}

object Demo24partitioner {
  def main(args: Array[String]): Unit = {

    val conf: SparkConf = new SparkConf()
      .setMaster("local[8]")
      .setAppName("pr")


    val sc = new SparkContext(conf)

    val rdd: RDD[String] = sc.textFile("data/students.txt")

    val kvRDD: RDD[(String, Int)] = rdd.map(line => (line.split(",")(4), 1))


    //使用自定义分区
    val countRDD: RDD[(String, Int)] = kvRDD.reduceByKey(new MyPartition, _ + _)

    countRDD.saveAsTextFile("data/partition")

  }

}

/**
  * 自定义分区
  *
  */
class MyPartition extends Partitioner {
  //设置分区数据
  override def numPartitions: Int = 3

  //获取key 属于哪一个分区
  override def getPartition(key: Any): Int = {
    val str: String = key.toString.substring(0, 2)

    if (str == "文科") {
      0
    } else {
      1
    }

  }

}
